import pandas as pd
from qlib.backtest.decision import Order, EmptyTradeDecision, TradeDecisionWO
from qlib.strategy.base import BaseStrategy
from qlib.contrib.strategy.signal_strategy import BaseSignalStrategy, TopkDropoutStrategy
from qlib.backtest.position import Position
from typing import Dict, List, Text, Tuple, Union, Any
import os
import copy

from qlib.backtest.decision import OrderDir
from qlib.backtest.signal import create_signal_from
import numpy as np
from scipy.optimize import bisect


class Strategy4SingleStock(BaseSignalStrategy):
    def __init__(
            self,
            *,
            change_freq,
            **kwargs,
    ):
        super().__init__(**kwargs)
        self.change_freq = change_freq

    def generate_trade_decision(
            self,
            execute_result: list = None,
    ):
        trade_step = self.trade_calendar.get_trade_step()
        if trade_step % self.change_freq != 1:
            return EmptyTradeDecision(self)
        trade_start_time, trade_end_time = self.trade_calendar.get_step_time(trade_step)
        pred_start_time, pred_end_time = self.trade_calendar.get_step_time(trade_step, shift=1)
        pred_score = self.signal.get_signal(start_time=pred_start_time, end_time=pred_end_time)
        current_temp: Position = copy.deepcopy(self.trade_position)
        code = pred_score.index[0]

        order_list = []
        cash = current_temp.get_cash()
        value = cash * self.risk_degree
        if pred_score.values[0] > 0:
            buy_price = self.trade_exchange.get_deal_price(
                stock_id=code, start_time=trade_start_time, end_time=trade_end_time, direction=OrderDir.BUY
            )
            buy_amount = value / buy_price
            factor = self.trade_exchange.get_factor(stock_id=code, start_time=trade_start_time, end_time=trade_end_time)
            buy_amount = self.trade_exchange.round_amount_by_trade_unit(buy_amount, factor)
            order = Order(
                stock_id=code,
                amount=buy_amount,
                start_time=trade_start_time,
                end_time=trade_end_time,
                direction=Order.BUY,  # 0 for sell, 1 for buy
            )
        else:
            stock_amount = current_temp.get_stock_amount(code=code)
            order = Order(
                stock_id=code,
                amount=stock_amount,
                start_time=trade_start_time,
                end_time=trade_end_time,
                direction=Order.SELL,  # 0 for sell, 1 for buy
            )
        order_list.append(order)
        return TradeDecisionWO(order_list, self)


class TimingStrategyOverTopKDropout(TopkDropoutStrategy):
    def __init__(self,
                 *,
                 time_signal,
                 change_freq=5,
                 topk=50,
                 n_drop=5,
                 **kwargs):
        super().__init__(topk=topk, n_drop=n_drop, **kwargs)
        time_signal: pd.DataFrame
        time_signal[:] = projection(time_signal.values, 0, 1)
        self.time_signal = create_signal_from(time_signal)
        self.change_freq = change_freq

    def generate_trade_decision(self, execute_result=None):
        trade_step = self.trade_calendar.get_trade_step()
        pred_start_time, pred_end_time = self.trade_calendar.get_step_time(trade_step, shift=1)
        pred_score = self.time_signal.get_signal(start_time=pred_start_time, end_time=pred_end_time)
        if trade_step % self.change_freq == 1:
            self.risk_degree = pred_score.values[0]
        return super().generate_trade_decision()


def sigmoid(x, k):
    return 1 / (1 + np.exp(-k * x))


def find_k(std, target_coverage=0.7) -> float:
    def coverage_difference(_k):
        s_m_minus_std = sigmoid(- std, _k)
        s_m_plus_std = sigmoid(std, _k)
        coverage = s_m_plus_std - s_m_minus_std
        return coverage - target_coverage

    k: float = bisect(coverage_difference, 0, 50)
    return k


def projection(data, a, b, target_coverage=0.7):
    mean = data.mean()
    std = data.std()
    data = data - mean
    data = data / std
    std = data.std()
    k = find_k(std, target_coverage)
    # 使用 Sigmoid 函数将原始数据映射到 [0, 1] 区间
    mapped_data_0_1 = sigmoid(data, k)

    # 将映射到 [0, 1] 区间的数据线性缩放到 [a, b] 区间
    mapped_data = a + (b - a) * mapped_data_0_1

    return mapped_data


if __name__ == "__main__":
    print(find_k(0.0017260506))
